#根据阈值线和地震目录绘制预报图
import py_tool
import py_stationsel
from datetime import datetime, timedelta
import scipy.io as sio
import py_pic
import argparse

import pandas as pd
import os

import argparse
import re

def is_float(s):
    pattern = r'^[+-]?\d*\.?\d+(?:[eE][+-]?\d+)?$'
    return bool(re.match(pattern,s))


def outputPredictDisp(iy, im, eqid, dist4sel=400, rank4sel=5.5, dayspan=360):
   #iy          = 2025
   #im          = 1
   #dist4sel    = 400
   #rank4sel    = 5.5
   #eqid        = 'X212MGPH0096'
   #yuzhi       = 0.122495714285714
   #days4span   = 360
   eqinfo      = "/home/pierce/project/predict4ts/速报目录.xls"
   folder_path =  "/home/pierce/project/predict4ts/TS4CON/";
   oroot_path  =  "/home/pierce/project/predict4ts/pred4pic/"
   eqroot_path =  "/home/pierce/project/predict4ts/{}_EQ.csv".format(eqid)
   py_tool.mkdir4file(oroot_path)
   oroot       =  "{}{:04}{:02}/".format(oroot_path, iy, im)
   py_tool.mkdir4file(oroot)
   STR_SDATE,STR_EDATE =  py_tool.getmonthdatainfo(iy,im,1)

   oblat,oblon         = py_stationsel.selposition4eqid(eqid)
   #if os.path.exists(eqroot_path):
   py_stationsel.seleqinfo(eqid, oblat, oblon, dist4sel, rank4sel, eqinfo, eqroot_path)
   NO, EQDATE, EQRANK, EQDIS, EQPOS = py_stationsel.seleqinfo4somecol(eqroot_path)
   #print('地震震级')
   #print(EQRANK)
   EQDATE = py_tool.str2matlabdate_list(EQDATE)
   EQPOS  = py_tool.transferlongname2shortname(EQPOS)
   #读取地震的编号, 震中距, 震级3个
   #print(EQPOS)
   keyword2    =  STR_EDATE
   keyword1    =  eqid
   ifiles      = py_tool.search_files_by_keyword_two(folder_path, keyword1, keyword2)
   if len(ifiles):
         for ifile in ifiles:
            mat_data = sio.loadmat(ifile)
            DDATE = mat_data['DDATE']
            #print(DDATE[0,0])
            DDATA = mat_data['DDATA']
            #print(DDATA.shape)
            #print(DDATA[0,0])
            DERRO = mat_data['DERRO']
            #print(DERRO.shape)
            MDATE    = mat_data['MDATE'] #根据它来筛选地震
            MVALUE   = mat_data['MVALUE']
            TimeSpan = mat_data['TimeSpan']
            EQDATE, EQRANK, EQDIS, EQPOS = py_tool.selTimeSpanEqContent(MDATE, EQDATE, EQRANK, EQDIS, EQPOS)
            iroot,isfile   = py_tool.getfileinfo(ifile)
            stnm,eqid,date = py_tool.splitfilename(isfile)
            ofile = "{}_{}_{}_A-PRED.png".format(stnm,eqid,date)
            py_pic.paintdisp4predictlongday(DDATE,  DDATA,    DERRO,    MDATE,
                                            MVALUE, TimeSpan, oroot,    ofile,
                                            yuzhi,  NO,       EQDATE,   EQRANK, 
                                            EQDIS,  EQPOS,    rank4sel, days4span, 
                                            stnm)
            ofile = "{}{}_{}_{}_A-PRED.txt".format(oroot, stnm,eqid,date)
            py_tool.outputpredictdescripsion(DDATE,  DDATA,  DERRO, MDATE, MVALUE, TimeSpan, 
                             EQDATE, EQRANK, EQDIS, EQPOS,
                             yuzhi,  days4span,
                             dist4sel, rank4sel,
                             ofile)
#根据阈值线和地震目录绘制预报图
if __name__ == "__main__":
   #iy          = 2025
   #im          = 1
   #dist4sel    = 400
   #rank4sel    = 5.5
   #eqid        = 'X212MGPH0096'
   #yuzhi       = 0.122495714285714
   #days4span   = 360
   #--date 202501 --eqid X212MGPH0096 --yuzhi 0.122495714285714 --dist4sel 400 --rank4sel 5.5  --days4span 360
   parser = argparse.ArgumentParser()
   parser.add_argument(
        '--date',
        type=str,  
        help="Specify the year and month in the format 'YYYYMM'. Defaults to the current year and month."
   )
   parser.add_argument(
        '--eqid',
        type=str,  
        help="重力仪的仪器ID, 默认值为: X212MGPH0096."
   )
   parser.add_argument(
        '--yuzhi',
        type=str,  
        help="设置预报的阈值, 大于这个值的为异常, 小于这个的为正常 默认值: yuzhi = 0.122495714285714."
   )
   parser.add_argument(
        '--dist4sel',
        type=str,  
        help="震中距选择地震目录, 默认值为 400 ."
   )
   parser.add_argument(
        '--rank4sel',
        type=str,  
        help="地震级选择大小, 默认值为: 5.5 ."
   )
   parser.add_argument(
        '--days4span',
        type=str,  
        help="预报有效时长, 单位天. 默认值: 360 ."
   )
   #--date 202501 --eqid X212MGPH0096 --yuzhi 0.122495714285714 --dist4sel 400 --rank4sel 5.5  --days4span 360
   # 设置args的默认值和读取参数值
   args = parser.parse_args()
   if args.date:
       if not re.match(r"^\d{6}$", args.date):
          raise argparse.ArgumentTypeError(f"Invalid format: '{args.date}'. Expected 'YYYYMM'.")
       else:
          year, month = args.date[:4], args.date[4:]
   else:
       new_data = datetime.datetime.today() - relativedelta(months=1)
       year, month = new_data.year, new_data.month 
   timeline = f"{year}年{str(month).rjust(2, '0')}月"
   iy    = int(year)
   im    = int(month)

   if args.eqid:
       if not re.match(r"^[A-Z0-9]+$", args.eqid):
          raise argparse.ArgumentTypeError(f"Invalid format: '{args.eqid}'. Expected 'X212MGPH0096'.")
       else:
          str_eqid = args.eqid
   else:
       str_eqid = 'X212MGPH0096'
   # 设置args的默认值和读取参数值
   eqid = str_eqid

   if args.yuzhi:
       if not is_float(args.yuzhi):
          raise argparse.ArgumentTypeError(f"Invalid format: '{args.yuzhi}'. Expected 0.1224.")
       else:
          str_yuzhi = args.yuzhi
   else:
       str_yuzhi = '0.122495714285714'
   # 设置args的默认值和读取参数值
   yuzhi = float(str_yuzhi)

   if args.dist4sel:
       if not args.dist4sel.isdecimal():
          raise argparse.ArgumentTypeError(f"Invalid format: '{args.dist4sel}'. Expected 400.")
       else:
          str_dist4sel = args.dist4sel
   else:
       str_dist4sel = '400'
   # 设置args的默认值和读取参数值
   dist4sel = float(str_dist4sel)

   if args.dist4sel:
       if not is_float(args.dist4sel):
          raise argparse.ArgumentTypeError(f"Invalid format: '{args.dist4sel}'. Expected 400.")
       else:
          str_dist4sel = args.dist4sel
   else:
       str_dist4sel = '400'
   # 设置args的默认值和读取参数值
   dist4sel = float(str_dist4sel)

   if args.rank4sel:
       if not is_float(args.rank4sel):
          raise argparse.ArgumentTypeError(f"Invalid format: '{args.rank4sel}'. Expected 5.5.")
       else:
          str_rank4sel = args.rank4sel
   else:
       str_rank4sel = '5.5'
   # 设置args的默认值和读取参数值
   rank4sel = float(str_rank4sel)

   if args.days4span:
       if not is_float(args.days4span):
          raise argparse.ArgumentTypeError(f"Invalid format: '{args.days4span}'. Expected 360.")
       else:
          str_days4span = args.days4span
   else:
       str_days4span = '360'
   # 设置args的默认值和读取参数值
   days4span = float(str_days4span)
   outputPredictDisp(iy, im, eqid, dist4sel, rank4sel, days4span)



